#include "StdAfx.h"
#include ".\dnsqueryworker.h"


int DnsQueryExample(void)
{


	char szResult[64] = {0};
	CDnsQueryWorker query;

	unsigned short id = (unsigned short) GetTickCount() % 0xFFFF;

	query.QueryDns( "geo.bnbt.net", "8.8.8.8", 0x0123 , DNS_TYPE_A );

	int i = 0;
	for ( i = 0 ; i < 20 ; i++ )
	{
		if( query.Update() != DNS_WORKER_STATUS_QUERY )
			break;

		Sleep( 100 );
	}

	strcpy( szResult, query.GetAddress() );

	return 0;
}



CDnsQueryWorker::CDnsQueryWorker(void)
{
	 sckClient = INVALID_SOCKET;
	 dwLast = 0;
	 m_bQuerying = false;

	 m_CacheDomain[0] = '\0';

	  nphp = new DNS_PKG_HEADER;  
	  dkqp = new DNS_PKG_QUERY;  

	  m_bStart = true;
}

CDnsQueryWorker::~CDnsQueryWorker(void)
{
	CloseWorker();
	delete nphp;
	delete dkqp;

}

int CDnsQueryWorker::QueryDns(const char* domain, const char* ip, unsigned short id, unsigned short dns_type, u_short dns_port )
{

	if( m_bStart )
	{
		m_bStart = false;
		Sleep( 200 );
	}

	dwLast = GetTickCount();

	m_dns_type = dns_type;

	OpenWorker();

	 m_CacheDomain[0] = '\0';
	strcpy( m_CacheDomain, domain );

	//nphp = new DNS_PKG_HEADER;  
	//dkqp = new DNS_PKG_QUERY;  
	unsigned char dnsBuff[1024] = {0};  
	unsigned char dnsRecv[2048] = {0};  
	ZeroMemory(dnsBuff, 1024);  
	ZeroMemory(dnsRecv, 2048);  
	
	m_send_size = pkgDNS(domain, nphp, dkqp, dnsBuff, id);//fulfil dns package  
	if(m_send_size == -1)  
		return false;  

	struct sockaddr_in server_addr;  
	unsigned int pkg_len = 0;  
	ZeroMemory(&server_addr, sizeof(struct sockaddr_in));  
	
	server_addr.sin_family = AF_INET;  	
	server_addr.sin_port = htons(dns_port); // htons(53);  
	server_addr.sin_addr.S_un.S_addr = inet_addr( ip );
	pkg_len = sizeof(server_addr);  
	int in_len = sendto(sckClient, (const char*)dnsBuff, m_send_size, 0,  
		(struct sockaddr*) &server_addr, pkg_len);  
	if (in_len < 0) {  
//		std::cout << "tiny dns:ERROR sendto" << strerror(errno)<<std::endl;  
		CloseWorker();
		return false;  
	}  

	m_bQuerying = true;

	return 0;
}

int CDnsQueryWorker::CloseWorker(void)
{
	m_bQuerying = false;
	if( sckClient != INVALID_SOCKET )
	{
		closesocket( sckClient );
		sckClient = INVALID_SOCKET;
	}

	
	return 0;
}

BOOL CDnsQueryWorker::OpenWorker(void)
{
	sckClient = socket( AF_INET, SOCK_DGRAM, 0 );
	memset( m_szIpAddress, 0, sizeof(m_szIpAddress) );
	return TRUE;
}

int CDnsQueryWorker::pkgDNS(const char* domain, DNS_PKG_HEADER_PTR& nphp,  
						 DNS_PKG_QUERY_PTR& dkqp, unsigned char* dnsBuff, unsigned short id)
{

	char tmpBuf[256];  
	ZeroMemory(tmpBuf, 256);  
	int domainLen = strlen( domain);  
	if(domainLen <= 0)  
		return -1;  
	memcpy(tmpBuf, domain, domainLen);  
	dkqp->dns_name = new unsigned char[domainLen + 2* sizeof (unsigned char)];  
	ZeroMemory(dkqp->dns_name, domainLen+2* sizeof (unsigned char));  
	char* tok = NULL;  
	tok = strtok(tmpBuf, ".");  
	unsigned char dot = '\0';  
	int offset = 0;  
	while(tok != NULL) {  
		dot = (unsigned char)strlen(tok);  
		memcpy(dkqp->dns_name +offset, &dot, sizeof(unsigned char));  
		offset += sizeof(unsigned char);  
		memcpy(dkqp->dns_name +offset, tok, strlen(tok));  
		offset += strlen(tok);  
		tok = strtok(NULL, ".");  
	}  
	dkqp->dns_name[domainLen+ 2*sizeof(unsigned char) -1] = 0x00;  
	nphp->id = htons(id);//dns transaction id, given randomly  
	nphp->flag = htons(0x0100); //dns standard query;  
	nphp->questions = htons(0x0001); //num of questions;  
	nphp->answerRRs = htons(0x0000);  
	nphp->authorityRRs = htons(0x0000);  
	nphp->additionalRRs = htons(0x0000);  
	//dkqp->dns_type = htons(0x0001); //Type  : A  
	//dkqp->dns_type = htons(0x0002); //Type  : A  
    dkqp->dns_type = htons( m_dns_type );
	dkqp->dns_class = htons(0x0001); //Class : IN  
	memcpy(dnsBuff, (unsigned char*)nphp, sizeof(DNS_PKG_HEADER));  
	memcpy(dnsBuff+sizeof(DNS_PKG_HEADER), (unsigned char*)dkqp->dns_name,domainLen+2*sizeof(unsigned char));  
	memcpy(dnsBuff+sizeof(DNS_PKG_HEADER)+(domainLen+2*sizeof(unsigned char)), (unsigned char*)&dkqp->dns_type, sizeof(unsigned short));  
	memcpy(dnsBuff+sizeof(DNS_PKG_HEADER)+(domainLen+2*sizeof(unsigned char))+sizeof(unsigned short), (unsigned char*)&dkqp->dns_class, sizeof(unsigned short));  
	delete[] dkqp->dns_name;  
	return sizeof(DNS_PKG_HEADER)+(domainLen+2*sizeof(unsigned char))+sizeof(unsigned short)+sizeof(unsigned short);  

	return 0;

}
int CDnsQueryWorker::Update( DWORD dwTimeout )
{
	
	WORD wList[20] = {0};

	//BYTE recvBuffer[4192] = {0};

	unsigned short flag = 0;

	memset( m_recvBuffer, 0, 20 );

	if( m_bQuerying )
	{
		    struct timeval tv;

			int nErrCode = 0;

			tv.tv_sec = dwTimeout / 1000;
			tv.tv_usec = dwTimeout % 1000;

			fd_set fdServer = {0};

			FD_ZERO( &fdServer );		

			FD_SET( sckClient, &fdServer );

			int nRet = select( 1, &fdServer, NULL, NULL, &tv );

			if( nRet == SOCKET_ERROR )
			{
				nErrCode = WSAGetLastError();
                CloseWorker();
				return DNS_WORKER_STATUS_FREE;
			}

			
			if( FD_ISSET( sckClient, &fdServer) )
			{
				
				struct sockaddr_in server_addr;  
				int pkg_len = sizeof(sockaddr_in);
				int nRet = recvfrom( sckClient, (char*)m_recvBuffer,  8190, 0,   (struct sockaddr*) &server_addr, &pkg_len); 
				
				
				memcpy( wList , m_recvBuffer, 16 );

				if( nRet < 0 )
				{
					nErrCode = WSAGetLastError();
					CloseWorker();
					return DNS_WORKER_STATUS_FREE;
				}
									

				unsigned char* p = (unsigned char*)m_recvBuffer;

				memset( m_szIpAddress, 0, sizeof(m_szIpAddress) );
				if( m_dns_type == DNS_TYPE_A ) // A Record
				{
				   bool analyseResult = recvAnalyse( m_recvBuffer, nRet, m_send_size, m_szIpAddress );
				   if( analyseResult )
				   {
					   return DNS_WORKER_STATUS_EXISTS;
				   }


				}

				p += 2;//dns id  
				flag =ntohs(*((unsigned short*) p));// p[0] * 0x100 + p[1];  
				if (flag != 0x8180) { 
				}

				
				int nStatus = DNS_WORKER_STATUS_EXISTS;								
				
				if( flag == 0x8183 || flag == 0x8381 )
					nStatus = DNS_WORKER_STATUS_NOT_EXISTS;
				else
					nStatus = DNS_WORKER_STATUS_EXISTS;


				CloseWorker();
			
				return nStatus;
			}
		

			DWORD dwNow = GetTickCount();
			if( dwNow > dwLast + (1000*10) )
			{
				CloseWorker();
				return DNS_WORKER_STATUS_FREE;
			}

			return DNS_WORKER_STATUS_QUERY;

	}

	return DNS_WORKER_STATUS_FREE;

	return 0;
}

char* CDnsQueryWorker::GetDomain(void)
{
	return m_CacheDomain;
}

char* CDnsQueryWorker::GetAddress(void)
{
	return m_szIpAddress;
}

bool CDnsQueryWorker::recvAnalyse(unsigned char* buf, size_t buf_size,    size_t send_size, char* ip) 
{  
									  bool result = false;  
									  unsigned char* p = buf;  
									  p += 2;//dns id  
									  unsigned short flag =ntohs(*((unsigned short*) p));// p[0] * 0x100 + p[1];  
									  if (flag != 0x8180) {  
										//  std::cout << "not a \"standard query response no error\"! "  
										//	  << std::endl;  
										  return false;  
									  }  
									  p += 2;//dns flag  
									  p += 2;//dns questions  
									  unsigned short answerRRs = ntohs(*((unsigned short*) p));//p[0] * 0x100 + p[1];//dns answer RRs  
									  p = buf + send_size;//p point to Answers  
									  unsigned short type;  
									  unsigned short dataLen;  
									  for (int i = 0; i < answerRRs; i++) {  
										  p += 2;//Name  
										  type = ntohs(*((unsigned short*) p));//p[0] * 0x100 + p[1];  
										  p += 2;//Type;  
										  if (type == 0x0001) {  
											  p += 2;//Class  
											  p += 4;//TTL  
											  p += 2;//Data Length  
											  strncpy(ip, inet_ntoa(*(struct in_addr*) p), 16);  
											  result = true;  
											  break;  
										  }  
										  p += 2;//Class  
										  p += 4;//TTL  
										  dataLen = ntohs(*((unsigned short*) p));//p[0] * 0x100 + p[1];  
										  p += 2;//Data Length  
										  p += dataLen;//data  
									  }  
									  return result;  
								  }  
